# -*- coding:utf-8 -*-
import asyncio
import logging
# @Time    : 2023/3/9 14:29
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : app.py
# @Software: chatbot

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import sys
import time
import traceback
from common.log import logger

from flask import Flask, render_template, request, Response, make_response, stream_with_context
from flask_cors import CORS

import config
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

arg = sys.argv
app = Flask(__name__)
CORS(app, resources=r'/*')
app.static_folder = 'static'


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/document')
def document():
    return render_template('index_document.html')


@app.route('/openai')
def openai():
    return render_template('index_openai.html')


import chat

config.load_config()
chat_bot = chat.Chat()
chat_bot.init("insurancePensionSteam", "123456")


# 页面初始化时，传入不同场景的参数，init后台
@app.route('/init')
def init():
    user_id = request.args.get('userId')
    user_name = request.args.get('userName')
    init_scene = request.args.get('initScene')
    reply = chat_bot.init(init_scene, user_id)

    # 存储用户和会话id对应关系

    return reply


@app.route('/user_info')
def user_info():
    session_id = request.args.get('sessionId')
    reply = chat_bot.get_user_info(session_id)
    return reply


@app.route('/message')
def message():
    message = request.args.get('message')
    user_id = request.args.get('userId')
    reply = chat_bot.get_reply(message, user_id)

    return reply


async def test(rush={}):
    total = 0
    for i in range(10):
        print(i)
        yield bytes(i)
        time.sleep(1)
        total += 1
    rush['total'] = total
    print(rush)

@app.route('/message_by_stream')
async def message_by_stream():
    message = request.args.get('message')
    user_id = request.args.get('userId')

    @stream_with_context
    def get_reply_by_stream():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        #tasks_1 = [chat_bot.get_reply_by_stream(message, user_id)]
        #loops = asyncio.get_event_loop()
        chat_reply = loop.run_until_complete(chat_bot.get_reply_by_stream(message, user_id))
        res_generator = chat_reply

        for res in res_generator:
            print(res)
            yield res

    response = Response(get_reply_by_stream, content_type='application/octet-stream')
    return response


@app.route('/document_assist')
def document_assist():
    document_chat_bot = chat.documentChat()
    message = request.args.get('message')
    user_id = request.args.get('userId')
    reply = document_chat_bot.get_reply(message, user_id)
    return reply


@app.route('/openai_func')
def openai_func():
    openai_chat_bot = chat.openaiChat()
    message = request.args.get('message')
    user_id = request.args.get('userId')
    reply = openai_chat_bot.get_reply(message, user_id)
    return reply


# 传入prompt，返回response
@app.route('/get_prompt_response')
def get_prompt_response():
    prompt = request.args
    reply = chat_bot.get_reply_with_prompt(prompt)
    return reply


if __name__ == '__main__':
    config.load_config()
    app.run(host='0.0.0.0', port=8080)
